from typing import Dict,Literal
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from xq_factor.factor import AbstractFactor
from xq_factor.extension import AbstractAnalyzer
import warnings
warnings.filterwarnings("ignore")

class HoldingAnalyzer(AbstractAnalyzer):
    def __init__(
        self,
        name: str=None,        
        benchmark_weights: AbstractFactor=None,
        industry: AbstractFactor=None,
        riskfactors: Dict[str,AbstractFactor]=None,
        config = None
    ):
        super().__init__(config)
        self.benchmark_weights = benchmark_weights
        self.industry = industry
        self.riskfactors = riskfactors
        self.name = name
    
    def process(self, market_value:AbstractFactor):
        market_value = self._get_value(market_value)
        benchmark_weights = self._get_value(self.benchmark_weights)
        industry = self._get_value(self.industry)
        if self.riskfactors is None:
            riskfactors = None
        else:
            riskfactors = {name:self._get_value(riskfactor) for name,riskfactor in self.riskfactors.items()}
        return HoldingAnalysisResult(self.name,market_value, benchmark_weights, industry,riskfactors)

class HoldingAnalysisResult:
    def __init__(self, name:str,market_value:pd.DataFrame, benchmark_weights:pd.DataFrame, industry,riskfactors:Dict[str,pd.DataFrame]):
        self.name = name
        self.market_value = market_value
        self.weights = market_value.div(market_value.sum(axis=1), axis=0)
        self.benchmark_weights = benchmark_weights
        self.industry = industry
        self.riskfactors = riskfactors
        self.industry_concentration = self.calc_industry_concentration()
        self.benchmark_industry_concentration = self.calc_industry_concentration(benchmark=True)
        self.riskfactor_exposure = self.calc_riskfactor_exposure()
        self.benchmark_riskfactor_exposure = self.calc_riskfactor_exposure(benchmark=True)
    
    def holding_num(self):
        return (self.market_value>0).sum(axis=1)

    def create_holding_num_figure(self):
        fig = px.line(self.holding_num())

        fig.update_layout(
            title='持仓数量随时间变化图',
            xaxis_title='时间',
            yaxis_title='持仓数量',
            legend_title='产品'
        )
        return fig
    
    def holdings_detail(self,top_n: int=None,date:pd.Timestamp=None):
        
        holdings = self.market_value.stack().to_frame()
        holdings.columns=["market_value"]
        holdings.sort_values(by=["datetime","market_value"],ascending=[True,False],inplace=True)
        holdings["weight"] = self.weights.stack()
        # # 将 '持仓权重' 列转换为百分比形式
        # holdings['position_weight'] = holdings['position_weight'].apply(lambda x: '{:.2%}'.format(x))

        if top_n is not None:
            weight_rank = self.market_value.rank(axis=1,ascending=False).stack()
            index = weight_rank[weight_rank<=top_n].index
            holdings = holdings.loc[index]
        if date is not None:
            date = pd.Timestamp(date)
            holdings = holdings.loc[date]
        return holdings

    def calc_industry_concentration(self,benchmark:bool=False)->pd.DataFrame:
        if self.industry is None:
            return None
        if benchmark:
            weights = self.benchmark_weights
        else:
            weights = self.weights
        df = pd.concat(
            [weights.stack(), self.industry.stack()], axis=1
        ).dropna()
        df.columns = ["factor", "grouper"]
        ind_df = (
            df.groupby(["datetime", "grouper"], group_keys=False)["factor"]
            .sum()
            .unstack(level="grouper")
        )
            # 将 '持仓权重' 列转换为百分比形式
        # ind_df = ind_df.applymap(lambda x: '{:.2%}'.format(x))
        return ind_df

    def calc_riskfactor_exposure(self,benchmark:bool=False)->pd.DataFrame:
        if self.riskfactors is None:
            return None
        if benchmark:
            weights = self.benchmark_weights
        else:
            weights = self.weights
        factors = []
        for risk_factor_name,risk_factor_data in self.riskfactors.items():
            weighted_exposure = weights * risk_factor_data
            weighted_exposure = weighted_exposure.sum(axis=1)
            weighted_exposure.name = risk_factor_name
            factors.append(weighted_exposure)
        return pd.concat(factors, axis=1)
    
    def create_industry_figure(self,mode:str=Literal["TS","CS"],param:pd.Timestamp|str=None):

        product_concentration = self.industry_concentration
        benchmark_concentration = self.benchmark_industry_concentration
        deviation = product_concentration - benchmark_concentration
        df = pd.concat([product_concentration,benchmark_concentration,deviation],axis=1,keys=['产品权重','基准权重','偏差'])
        df = df.swaplevel(0,1,axis=1).sort_index(axis=1)
        if mode == "TS":
            ts_df = df[param]

            # 绘制折线图（产品权重和基准权重）
            fig = px.line(ts_df, x=ts_df.index.astype(str), y=['产品权重', '基准权重'])

            # 设置左轴为默认的 y 轴
            fig.update_traces(yaxis="y")

            # 添加柱状图（偏差）
            fig.add_bar(x=ts_df.index.astype(str), y=ts_df['偏差'], name='偏差', yaxis="y2")

            # 更新布局
            fig.update_layout(
                title='产品权重、基准权重和偏差',
                xaxis=dict(
                    type='category',  # 设置 x 轴为分类类型
                    title='日期'
                ),
                yaxis=dict(
                    title='权重',
                    side="left",
                    showgrid=False
                ),
                yaxis2=dict(
                    title='偏差',
                    side="right",
                    overlaying="y",
                    showgrid=False
                ),
                legend_title='数据类型',
                bargap=0.3  # 设置不同 index 之间的柱状图间距，取值范围 0-1，值越大间距越大
            )
        else:
            cs_df = df.loc[param].unstack()
            fig = px.bar(cs_df, x=cs_df.index, y=['产品权重', '基准权重'],barmode='group')


            # 添加折线图（偏差）
            fig.add_scatter(x=cs_df.index, y=cs_df['偏差'], mode='lines', name='偏差')

            # 更新布局
            fig.update_layout(
                title='产品权重、基准权重和偏差',
                xaxis_title='日期',
                yaxis_title='权重/偏差',
                legend_title='数据类型',
                bargap=0.3  # 设置不同 index 之间的柱状图间距，取值范围 0-1，值越大间距越大
            )            
        return fig

    def create_riskfactors_figure(self,mode:str=Literal["TS","CS"],param:pd.Timestamp|str=None,bound_std=0.5):

        product_exposure = self.riskfactor_exposure
        benchmark_exposure = self.benchmark_riskfactor_exposure
        active_exposure = product_exposure-benchmark_exposure
        # 计算指定标准差的上下轨
        upper_bound = benchmark_exposure+bound_std*benchmark_exposure.std()
        lower_bound = benchmark_exposure-bound_std*benchmark_exposure.std()       

        df = pd.concat([product_exposure,benchmark_exposure,active_exposure,upper_bound,lower_bound],axis=1,keys=['产品暴露','基准暴露',"主动暴露",'上轨',"下轨"])
        df = df.swaplevel(0,1,axis=1).sort_index(axis=1)

        if mode == "TS":
            ts_df = df[param]

            # 绘制柱状图（产品权重和基准权重）
            fig = go.Figure()

            # 添加下轨折线
            fig.add_trace(
                go.Scatter(
                    x=ts_df.index,
                    y=ts_df["下轨"],
                    mode='lines',
                    line=dict(color='blue', width=0.5),
                    name='下轨',
                    showlegend=False
                )
            )

            # 添加填充区域
            fig.add_trace(
                go.Scatter(
                    x=ts_df.index,
                    y=ts_df["上轨"],
                    mode='lines',
                    line=dict(color='rgba(0,0,0,0)'),
                    fillcolor='rgba(0, 0, 255, 0.2)',  # 透明蓝色填充
                    fill='tonexty',
                    name='上下轨区间',
                    showlegend=False
                )
            )

            # 添加上轨折线
            fig.add_trace(
                go.Scatter(
                    x=ts_df.index,
                    y=ts_df["上轨"],
                    mode='lines',
                    line=dict(color='blue', width=0.5),
                    name='上轨',
                    showlegend=False
                )
            )

            # 添加折线图（偏差）
            fig.add_scatter(x=ts_df.index, y=ts_df['产品暴露'], mode='lines', name='产品暴露')

            # 更新布局
            fig.update_layout(
                title='产品权重、基准权重和偏差',
                xaxis_title='风格因子',
                yaxis_title='权重/偏差',
                legend_title='数据类型',
                bargap=0.3  # 设置不同 index 之间的柱状图间距，取值范围 0-1，值越大间距越大
            )
        else:
            cs_df = df.loc[param].unstack()
            # 绘制柱状图（产品权重和基准权重）
            fig = px.bar(cs_df, x=cs_df.index, y=['产品暴露', '基准暴露'],barmode='group')


            # 添加折线图（偏差）
            fig.add_scatter(x=cs_df.index, y=cs_df['主动暴露'], mode='lines', name='主动暴露')

            # 更新布局
            fig.update_layout(
                title='产品权重、基准权重和偏差',
                xaxis_title='风格因子',
                yaxis_title='权重/偏差',
                legend_title='数据类型',
                bargap=0.3  # 设置不同 index 之间的柱状图间距，取值范围 0-1，值越大间距越大
            )
        return fig
        